package com.hust.yuqian;

import com.hust.yuqian.attacks.NoiseAdditionFilter;
import com.hust.yuqian.dcts.FDct;
import com.hust.yuqian.dcts.IDct;
import com.hust.yuqian.tests.PNSR;
import com.hust.yuqian.utils.*;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;

public class Main {

    public static void main(String[] args) throws Exception{
        String Lenna = "Lenna.jpg";
        String lennaPath = (new StringBuilder(System.getProperty("user.dir")))
                .append(System.getProperty("file.separator")).append(Lenna).toString();
        String info = "  余谦U201014640";
        String waterMarkedImgName = "LennaWaterMark.jpg";
        String extractWaterMarkedImgName = "extractWaterMark.jpg";
        String GaussionNoiseImgName = "GaussionNoiseImg.jpg";
        String PoissonNoiseImgName = "PoissonNoiseImg.jpg";

        ImageTools imageTools = new ImageTools();
        BufferedImage bufferedImage = imageTools.getImage(lennaPath);


        System.out.println("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~");
        System.out.println("请耐心等待10s, 循环嵌套太多, double太多,很费时间...");

        //打上水印
        double [][] dct = FDct.operation(UsefulData.image_height, UsefulData.image_width,
                imageTools.getImageB(bufferedImage));

        CDMASPSP cdmaspsp = new CDMASPSP();
        int [][] C = cdmaspsp.operation(info);

        System.out.printf("在当前目录下, 带有您 %s 信息的水印已经生成!\n", info);


        MatrixTools matrixTools = new MatrixTools();
        double [][] imbedded = matrixTools.Imbed(C, dct);

        int [][] b = IDct.operation(UsefulData.image_height, UsefulData.image_width, imbedded);

        BufferedImage waterMarkedImage = imageTools.changeImageRGB(bufferedImage, b);
        imageTools.WriteImage(waterMarkedImage, waterMarkedImgName);

        System.out.println("在当前目录下, 打上水印的图片已经生成!");

        //提取水印

        /**
         * 这里有一个已知的小 bug.
         * 见下面注释. 直接读取LennaWaterMark,再计算 b 矩阵
         * 会发现 b 矩阵很大问题.
         * bug 在 ImageIO.write(waterMarkedImage, "jpg", file)
         * 这一句产生.
         * 小 bug, 很奇怪. 不想解决了...
         */
        int [][] Cextract = matrixTools.Extract(FDct.operation(UsefulData.image_height, UsefulData.image_width,
                imageTools.getImageB(waterMarkedImage)), FDct.operation(UsefulData.image_height, UsefulData.image_width,
                imageTools.getImageB(imageTools.getImage(lennaPath))));
//        int [][] Cextract = matrixTools.Extract(FDct.operation(UsefulData.image_height, UsefulData.image_width,
//                imageTools.getImageB(ImageIO.read(new File("/Users/mmutuyu/IdeaProjects/watermark_cdma/LennaWaterMark.jpg")))),
//                FDct.operation(UsefulData.image_height, UsefulData.image_width,
//                        imageTools.getImageB(imageTools.getImage("lennaPath))));
        int [][] watermark_extract = cdmaspsp.operation_inverse(Cextract);

        BufferedImage extractImg = MakeWaterMarkPic.makePic(watermark_extract);
        imageTools.WriteImage(extractImg, extractWaterMarkedImgName);

        System.out.println("在当前目录下, 提取出来的水印图片已经生成!");
        System.out.println();
        System.out.println("测试部分:");

        //计算 PNSR
        PNSR.calculate(UsefulData.WATERMARK_HEIGHT, UsefulData.WATERMARK_WIDTH,
                imageTools.getImageB(imageTools.getImage(lennaPath)), imageTools.getImageB(waterMarkedImage));

        System.out.println();
        System.out.println("实验结束~~~");
        System.out.println("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~");
    }
}
